#
# Copyright (c) 2009, 2010, Davide Cittaro
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the <organization> nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

# This module introduces objects to handle
# spectra

import numpy as np
import scipy.stats as ss
from config import Elements
from zlib import crc32
from hashlib import md5
import os


def match(massArray, massToMatch, options):
  """ Returns the list of masses which
  are close enough to the given mass
  """
  tolerance = options["ms2tolerance"]
  unit = options["unit2"]
  if unit == 'ppm':
    tolerance = massToMatch / 1000000 * tolerance
  ret =  (massArray - tolerance <= massToMatch) & (massArray + tolerance >= massToMatch)
  return ret


class ms2spectrum:
  """ A Class which defines MS/MS spectra
  """
  def __init__(self):
    self._precursorMz = 0.0
    self._charge = 1
    self._description = ""
    self._masses = np.array([], dtype=float)
    self._intensities = np.array([], dtype=float)
    self._sorted = False
    self._fileName = ''
    self._checksum = None
    self._bestMatch = None
    self._matches = []
    self._pValue = 0.0
  def setFileName(self, file):
    self._fileName = file
  def fileName(self, fullPath = False):
    if fullPath:
      return os.path.abspath(self._fileName)
    return os.path.basename(self._fileName)
  def initWithList(self, aList):
    aList = [x for x in aList if len(x) > 0]
    tempSp = [x.split() for x in aList[4:]]
    self.setPrecursorMz(aList[1][8:])
    self.setCharge(int(aList[2][7:8]))
    self.setDescription(aList[3][6:])
    self._masses = np.array([float(x[0]) for x in tempSp])
    self._intensities = np.array([float(x[1]) for x in tempSp])
  def setSpectrum(self, tempList, axis = 1):
    """ The method is feeded with mass/intensities values. 
    if the array is
    [[m1, i1], [m2, i2]] 
    then axis should be 1
    if the array is
    [[m1, m2], [i1, i2]]
    then the axis should be 0
    """
    if axis == 1:
      self._masses = np.array([float(x[0]) for x in tempList])
      self._intensities = np.array([float(x[1]) for x in tempList])
    elif axis == 0:
      self._masses = np.array([float(x) for x in tempList[0]])
      self._intensities = np.array([float(x) for x in tempList[1]])
  def appendPeaks(self, tempList, axis = 1):
    """ The method is feeded with mass/intensities values. 
    if the array is
    [[m1, i1], [m2, i2]] 
    then axis should be 1
    if the array is
    [[m1, m2], [i1, i2]]
    then the axis should be 0
    """
    if axis == 1:
      self._masses = np.append(self._masses, np.array([float(x[0]) for x in tempList]))
      self._intensities = np.append(self._intensities, np.array([float(x[1]) for x in tempList]))
    elif axis == 0:
      self._masses = np.append(self._masses, np.array([float(x) for x in tempList[0]]))
      self._intensities = np.append(self._intensities, np.array([float(x) for x in tempList[1]]))
  def setPrecursorMz(self, value):
    self._precursorMz = float(value)
  def precursorMz(self, r = 6):
    return round(self._precursorMz, r)
  def setCharge(self, value):
    self._charge = int(value)
  def charge(self):
    return self._charge
  def precursorMass(self):
    return self._precursorMz * self._charge - (self._charge * Elements.getfloat("H", "mmass"))
  def length(self):
    return (len(self._masses))
  def setDescription(self, value):
    self._description = value.strip()
  def description(self):
    return self._description
  def addPeak(self, peak):
    if len(peak):
      self._masses = np.append(self._masses, float(peak[0]))
      self._intensities = np.append(self._intensities, float(peak[1]))
  def peaks(self):
    sp = np.array([self._masses, self._intensities]).transpose()
    return sp
  def peakMasses(self):
    return self._masses
  def peakIntensities(self):
    return self._intensities
  def maxIntensity(self):
    p = self._intensities.argmax()
    return self._intensities[p]
  def minIntensity(self):
    p = self._intensities.argmin()
    return self._intensities[p]
  def mostIntense(self, n = 1):
    mask = self._intensities.argsort()[-n:]
    mask.sort()
    return np.array([self._masses[mask], self._intensities[mask]])
  def slice(self, min = False, max = False):
    """ Returns a new m2spectra which contains data
    only for the interval specified
    """
    if not min:
      min = self._masses[0]
    if not max:
      max = self._masses[-1]
    mask = (self._masses >= min) & (self._masses <= max)
    r = ms2spectrum()
    r._masses = self._masses[mask]
    r._intensities = self._intensities[mask]
    r.setPrecursorMz(self._precursorMz)
    r.setCharge(self._charge)
    r.setDescription(self._description)
    return r
  def closeTo(self, mass, tolerance = 0.5, unit = "Da"):
    """ Returns a new m2spectrum which masses are close to the given
    mass
    """
    close = ms2spectrum()
    tolerance = float(tolerance)
    if unit == 'ppm':
      tolerance = mass / 1000000 * tolerance
    close.setPrecursorMz(self._precursorMz)
    close.setCharge(self._charge)
    close.setDescription(self._description)
    mask = np.abs(self._masses - mass) <= tolerance
    close._masses = self._masses[mask]
    close._intensities = self._intensities[mask]
    return close
  def closestTo(self, mass, tolerance = 0.5, unit = "Da"):
    """ Returns the closest peak to a certain mass
    """
    clsp = self.closeTo(mass, tolerance, unit) 
    if not len(clsp._masses):
      return np.array([])
    m = np.array(np.abs(clsp._masses - mass)).argmin()
    return np.array([clsp._masses[m], clsp._intensities[m]])
  def filterOnIntensity(self, threshold = 0):
    """ Filter peak data for a certain intensity
    Returns a new ms2spectrum() object
    """
    filtered = ms2spectrum()
    filtered.setPrecursorMz(self._precursorMz)
    filtered.setCharge(self._charge)
    filtered.setDescription(self._description)
    mask = self._intensities >= threshold
    filtered._masses = self._masses[mask]
    filtered._intensities = self._intensities[mask]
    return filtered
  def topIons(self, n = 12, window = 30):
    """ Filter peak data so that there are the most n intense
    peaks each m/z window spanning a given value.
    
    n: the number of most intense ion in the window
    window: the window width
    
    Returns a new ms2spectrum() object
    """
    filtered = ms2spectrum()
    filtered.setPrecursorMz(self._precursorMz)
    filtered.setCharge(self._charge)
    filtered.setDescription(self._description)
    filtered._checksum = self._checksum
    for x in np.arange(int(self._masses[0]), int(self._masses[-1]), window):
      filtered.appendPeaks(self.slice(x, x+window).mostIntense(n), axis = 0)
    return filtered
  def checksum(self):
    """ Returns the md5 checksum for this spectrum
    """
    if not self._checksum:
      self._checksum = md5(':'.join([str(x) for x in [self.description(), self.charge(), self.precursorMz(), self.length()]])).hexdigest()
    return self._checksum
  def span(self):
    """ Returns the mass span
    """
    return self._masses[-1] - self._masses[0]
  def setBestMatch(self, match):
      self.orderMatches()
  def bestMatch(self):
    return self._matches[0]
  def addMatch(self, m):
    """ Add a match with a peptide or xlink to this
    spectrum
    """
    csset = set([x.checksum() for x in self._matches])
    if m.checksum() not in csset:
      self._matches.append(m)
  def matches(self):
    return self._matches
  def removeMatch(self, match):
    newMatches = []
    for m in self._matches:
      if m.checksum() != match.checksum():
        newMatches.append(m)
    self._matches = newMatches
  def orderMatches(self):
    tmplist = [(x.pValue(self), x) for x in self._matches]
    tmplist.sort()
    self._matches = [x[1] for x in tmplist]
    
def randomSpectrum(sp, options):
  outSp = ms2spectrum()
  outSp.setPrecursorMz(sp.precursorMz())
  outSp.setCharge(sp.charge())
  outSp.setDescription(sp.description() + " random generated")

  mmean = np.average(sp._masses)
  mvar = np.var(sp._masses)
  imean = np.average(sp._intensities)
  ivar = np.var(sp._intensities)
# negative binomial
  mp = mmean / mvar
  mr = mmean * mp / (1 - mp)
  ip = imean / ivar
  ir = imean * ip / (1 - ip)
  for i in range(sp.length()):
    randMass = np.random.negative_binomial(mr, mp)
    tolerance = options["ms2tolerance"] * randMass / 1000000
    randMass += np.random.uniform(-tolerance, tolerance)
    randInt = np.random.negative_binomial(ir, ip) + np.random.uniform(-tolerance, tolerance)
    outSp.addPeak([randMass, randInt])
  return outSp  
